from setfit import AbsaModel, AbsaTrainer
from datasets import load_dataset

# You can initialize a AbsaModel using one or two SentenceTransformer models, or two ABSA models
# model = AbsaModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model = AbsaModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2", "sentence-transformers/all-mpnet-base-v2")

# The training/eval dataset must have `text`, `span`, `polarity`, and `ordinal` columns
raw_dataset = load_dataset("data", data_files="example_training_file.csv")
train_dataset = raw_dataset["train"].rename_columns({"sentence": "text", "aspect": "span", "polarity": "label"})

# The minimal Trainer instantiation
trainer = AbsaTrainer(model, train_dataset=train_dataset)
trainer.train()

from setfit import AbsaModel, AbsaTrainer, TrainingArguments
from transformers import EarlyStoppingCallback
from datasets import load_dataset

# You can initialize a AbsaModel using one or two SentenceTransformer models, or two ABSA models
# model = AbsaModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model = AbsaModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2", "sentence-transformers/all-mpnet-base-v2")

# The training/eval dataset must have `text`, `span`, `polarity`, and `ordinal` columns
raw_dataset = load_dataset("data", data_files="example_training_file.csv")["train"]
raw_dataset = raw_dataset.rename_columns({"sentence": "text", "aspect": "span", "polarity": "label"})
raw_dataset = raw_dataset.train_test_split(test_size=10)
train_dataset, eval_dataset = raw_dataset["train"], raw_dataset["test"]

# Training arguments for aspect and polarity training
aspect_args = TrainingArguments(
    output_dir="aspect",
    num_epochs=2,
    body_learning_rate=5e-5,
    head_learning_rate=1e-2,
    use_amp=True,
    warmup_proportion=0.2,
    evaluation_strategy="steps",
    eval_steps=20,
    save_steps=20,
    load_best_model_at_end=True,
)
polarity_args = TrainingArguments(
    output_dir="polarity",
    num_epochs=3,
    max_steps=1000,
    body_learning_rate=2e-5,
    head_learning_rate=3e-2,
    evaluation_strategy="steps",
    eval_steps=20,
    save_steps=20,
    load_best_model_at_end=True,
)

trainer = AbsaTrainer(
    model,
    args=aspect_args,
    polarity_args=polarity_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
)
trainer.train()

metrics = trainer.evaluate()
print(metrics)

# Predicting is as easy as with SetFit
predictions = model.predict([
    "Best pizza outside of Italy and really tasty.",
    "The variations are great and the prices are absolutely fair.",
    "Unfortunately, you have to expect some waiting time and get a note with a waiting number if it should be very full."
])
print(predictions)
"""
[
    [{'span': 'pizza', 'polarity': 'positive'}],
    [{'span': 'variations', 'polarity': 'positive'}, {'span': 'prices', 'polarity': 'positive'}],
    [{'span': 'waiting time', 'polarity': 'negative'}, {'span': 'note', 'polarity': 'positive'}, {'span': 'number', 'polarity': 'negative'}]
]
"""

# You can push to the Hub/save models using one or two repo_ids:
trainer.push_to_hub("tomaarsen/setfit-absa-restaurant-review", private=True)
# trainer.push_to_hub("tomaarsen/setfit-absa-restaurant-review-aspect", "tomaarsen/setfit-absa-restaurant-review-polarity", private=True)
# Or directly on the model:
# model.push_to_hub("tomaarsen/setfit-absa-restaurant-review-aspect", "tomaarsen/setfit-absa-restaurant-review-polarity", private=True)
# model.save_pretrained("absa-model")
# model.save_pretrained("absa-model-aspect", "absa-model-polarity", private=True)